package com.ideaaedi.component.compile;

import javax.tools.FileObject;
import javax.tools.ForwardingJavaFileManager;
import javax.tools.JavaFileManager;
import javax.tools.JavaFileObject;
import javax.tools.StandardLocation;
import java.io.IOException;
import java.io.InputStream;
import java.util.ArrayList;
import java.util.List;
import java.util.Map;
import java.util.Set;
import java.util.concurrent.ConcurrentHashMap;

/**
 * 自定义{@link JavaFileManager}文件管理器，用于管.java文件、.class文件
 * <p>
 *     为什么要重写{@link JavaFileManager}？
 * </p>
 *
 * @author JustryDeng
 * @since 2021/9/20 16:44:07
 */
public class DynamicJavaFileManager extends ForwardingJavaFileManager<JavaFileManager> {

    
    private final Map<String, CharSequenceJavaFileObject> javaFileObjectMap = new ConcurrentHashMap<>(16);
    
    private final DynamicClassLoader classLoader;
    
    private final PackageInternalsFinder finder;
    
    public DynamicJavaFileManager(JavaFileManager fileManager, DynamicClassLoader classLoader) {
        super(fileManager);
        this.classLoader = classLoader;
        this.finder = new PackageInternalsFinder(classLoader);
    }
    
    @Override
    public FileObject getFileForInput(Location location, String packageName, String relativeName) throws IOException {
        JavaFileObject javaFileObject = javaFileObjectMap.get(determineKey(location, packageName, relativeName));
        if (javaFileObject != null) {
            return javaFileObject;
        }
        return super.getFileForInput(location, packageName, relativeName);
    }
    
    /**
     * 这里是编译器返回的同(源)Java文件对象,替换为CharSequenceJavaFileObject实现
     */
    @Override
    public JavaFileObject getJavaFileForOutput(Location location, String classLongName, JavaFileObject.Kind kind, FileObject sibling) throws IOException {
        CharSequenceJavaFileObject javaFileObject = new CharSequenceJavaFileObject(classLongName, (InputStream) null);
        classLoader.registerCompiledSource(classLongName, javaFileObject);
        return javaFileObject;
    }
    
    @Override
    public ClassLoader getClassLoader(Location location) {
        // 使用自定义的类加载器
        return classLoader;
    }
    
    @Override
    public String inferBinaryName(Location location, JavaFileObject file) {
        if (file instanceof CharSequenceJavaFileObject || file instanceof ClassJavaFileObject) {
            return file.getName();
        }
        return super.inferBinaryName(location, file);
    }
    
    /**
     * 寻找JavaFileObject
     * <p>
     *     注：这个方法很关键，一个.java文件，可能依赖了很多其他依赖，此时，如果只对这个文件编译，那么会报错提示找不到对应的依赖，
     *         而通过这个方法就可以解决这个问题的
     * </p>
     */
    @Override
    public Iterable<JavaFileObject> list(Location location, String packageName, Set<JavaFileObject.Kind> kinds, boolean recurse) throws IOException {
        List<JavaFileObject> result = new ArrayList<>(16);

        // 先将原来的也加入进来
        for (JavaFileObject javaFileObject : super.list(location, packageName, kinds, recurse)) {
            result.add(javaFileObject);
        }
        
        // 这里要区分编译的Location以及编译的Kind
        if (location == StandardLocation.CLASS_PATH && kinds.contains(JavaFileObject.Kind.CLASS)) {
            // .class文件以及classPath下, 这里需要额外添加类加载器加载的JavaFileObject
            for (JavaFileObject file : classLoader.listCharSequenceJavaFileObject()) {
                if (file.getName().startsWith(packageName)) {
                    result.add(file);
                }
            }
            // 确保jar包中的依赖也能被解析到
            result.addAll(finder.find(packageName));
        } else if (location == StandardLocation.SOURCE_PATH && kinds.contains(JavaFileObject.Kind.SOURCE)) {
            // .java文件以及编译路径下
            for (JavaFileObject file : javaFileObjectMap.values()) {
                if (file.getName().startsWith(packageName)) {
                    result.add(file);
                }
            }
        }
        return result;
    }
    
    /**
     * 得到唯一key
     */
    private static String determineKey(Location location, String packageName, String relativeName) {
        return location.getName() + '/' + packageName + '/' + relativeName;
    }
}
